import os
import win32com.client
import re
import json
import sys
import shutil

'''
@Author:cosann
@email: 1066190125@qq.com
@Script_Name:Reshortcut.py
@Version:1.0
@Create_Date:20230309
@Description:
Windows快捷方式备份还原脚本
@Update content:
'''


def backup_src_shortcut(f_path):
    # 创建备份文件夹
    backup_dir = f_path + '\\shortcut_backup'
    if not os.path.exists(backup_dir):
        os.makedirs(backup_dir)

    print("backup_dir: ", backup_dir)
    for link_str in os.listdir(f_path):
        if link_str.endswith('.lnk'):
            full_shortcut = os.path.join(f_path, link_str)
            # print(full_shortcut)

            shutil.copy(full_shortcut, backup_dir)
    f_result_code = 0
    return backup_dir, f_result_code


def make_target_shortcut(f_path):
    # 重写path
    # f_path = "C:\\Users\\Cosann\\Desktop\\dirs"

    # 调用函数备份快捷方式
    back_dir, result_code = backup_src_shortcut(f_path)

    print("源快捷方式备份成功,备份路径: {}\n".format(back_dir)) if result_code == 0 else print("源快捷方式备份失败")
    # if result_code == 0:
    #     print("源快捷方式备份成功,备份路径: {}".format(back_dir))

    shell = win32com.client.Dispatch("WScript.Shell")

    shortcut_list = []
    for fn in os.listdir(f_path):

        if fn.endswith('.lnk'):
            com_path = f_path + '\\' + fn
            shortcut = shell.CreateShortCut(com_path)
            shortcut_dict = {fn: shortcut.Targetpath}
            # print(fn)
            shortcut_list.append(shortcut_dict)
            # print(shortcut.Targetpath)
    print(shortcut_list)

    # 将快捷方式名称和对应的真实路径写入json文件
    with open('Rebuild_shortcut.json', 'w', encoding='utf-8') as obj:
        # obj.write(str(re_shortcut_list))
        json.dump(shortcut_list, obj)
    # print(os.stat('Rebuild_shortcut.json'))
    if os.path.exists('Rebuild_shortcut.json') and os.stat('Rebuild_shortcut.json').st_size > 2:
        print(f"\n生成数据源文件Rebuild_shortcut.json成功\n")
    else:
        print(f"\n生成数据源文件失败,请检查!\n")


def del_src_shortcut(f_src_shortcut):
    # 删除源快捷方式
    # 遍历当前文件下的所有文件
    for src_shortcut in os.listdir(os.getcwd()):
        # 若是快捷方式执行
        if src_shortcut.endswith('.lnk'):
            # 若包含在字典中
            if src_shortcut in f_src_shortcut:
                # 删除
                os.remove(src_shortcut)


def create_shortcut(f_target_disk):
    f_target_disk = target_disk + ':'

    # os.chdir('E:\\test')
    shell = win32com.client.Dispatch("WScript.Shell")

    with open('Rebuild_shortcut.json', 'r', encoding='utf-8') as obj:
        shortcut_list = json.load(obj)

    # 修改真实路径盘符
    re_shortcut_list = []
    for shortcut_dict in shortcut_list:

        #调用删除函数
        del_src_shortcut(shortcut_dict)

        # print(shortcut_dict)
        for key, value in shortcut_dict.items():  # items方法返回一个键值对列表

            re_shortcut_dict = {key: re.sub('^[A-Z]:', f_target_disk, value)}
            # print(re_shortcut_dict)
            re_shortcut_list.append(re_shortcut_dict)

    print(re_shortcut_list)

    for shortcut_dict in re_shortcut_list:
        # print(shortcut_dict)
        for key, value in shortcut_dict.items():
            # print("key=", key)
            dest_shortcut = shell.CreateShortCut(key)
            dest_shortcut.TargetPath = value
            dest_shortcut.save()
            # print("key=", key)
            print("\n快捷方式[{}]创建成功!\n".format(key))


if __name__ == '__main__':

    # target_disk = 'K'
    if sys.argv[1:]:
        array_sum = len(sys.argv[1:])

        # 位置参数2为Backup
        if sys.argv[1] == "-b" and array_sum == 1:
            # print("backup")
            scripts_path = os.getcwd()
            print("scripts_path: " + scripts_path)
            # 创建目标字典
            make_target_shortcut(scripts_path)

        # 位置参数2为Make
        elif sys.argv[1] == "-m" and array_sum == 2:
            # print("make")
            target_disk = sys.argv[2]
            # 创建快捷方式
            create_shortcut(target_disk)

        else:
            script_name = sys.argv[0]
            print("脚本用法:\n1. 备份源快捷方式数据 命令:[ {} -b ]".format(
                script_name))
            print("2. 重新创建快捷方式 命令:[ {} -m 盘符 ]".format(script_name))
    else:
        script_name = sys.argv[0]
        print("脚本用法:\n1. 备份源快捷方式数据 命令:[ {} -b ]".format(
            script_name))
        print("2. 重新创建快捷方式 命令:[ {} -m 盘符 ]".format(script_name))